#include "vkop.h"
#include "vkutils.h"

namespace vkinfer
{
    VkOp::~VkOp()
    {
        if (shaderModule != VK_NULL_HANDLE)
            vkDestroyShaderModule(device, shaderModule, nullptr);
        if (descriptorPool != VK_NULL_HANDLE)
            vkDestroyDescriptorPool(device, descriptorPool, nullptr);
        if (descriptorSetLayout != VK_NULL_HANDLE)
            vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);
        if (pipelineLayout != VK_NULL_HANDLE)
            vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
        if (pipeline != VK_NULL_HANDLE)
            vkDestroyPipeline(device, pipeline, nullptr);
    }

    void VkOp::initialize(VkDevice device, const std::string& shader_file, uint32_t num_outputs, uint32_t num_inputs, 
        uint32_t num_weights, uint32_t pushconstants_size)
    {
        this->device = device;

        uint32_t num_descriptors = num_outputs + num_inputs + num_weights;
        create_descriptor_set_layout(num_descriptors);
        create_descriptor_set(num_descriptors);
        load_shader(shader_file);
        create_pipeline(pushconstants_size);
    }

    void VkOp::set_context(const std::shared_ptr<VkCompute>& cmd, const std::shared_ptr<VkTransfer>& transfer)
    {
        this->cmd = cmd;
        this->transfer = transfer;
    }

    void VkOp::create_descriptor_set_layout(uint32_t num_descriptors)
    {
        std::vector<VkDescriptorSetLayoutBinding> bindings(num_descriptors);

        for (uint32_t i = 0; i < num_descriptors; ++i)
        {
            bindings[i] = {};
            bindings[i].binding = i;
            bindings[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
            bindings[i].descriptorCount = 1;
            bindings[i].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
        }

        VkDescriptorSetLayoutCreateInfo descriptorSetLayoutCreateInfo = {};
        descriptorSetLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
        descriptorSetLayoutCreateInfo.bindingCount = (uint32_t)bindings.size();
        descriptorSetLayoutCreateInfo.pBindings = bindings.data();

        VK_CHECK_RESULT(vkCreateDescriptorSetLayout(device, &descriptorSetLayoutCreateInfo, NULL, &descriptorSetLayout));
    }

    void VkOp::create_descriptor_set(uint32_t num_descriptors)
    {
        VkDescriptorPoolSize descriptorPoolSize = {};
        descriptorPoolSize.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
        uint32_t num_bindings = num_descriptors;
        descriptorPoolSize.descriptorCount = num_bindings;

        VkDescriptorPoolCreateInfo descriptorPoolCreateInfo = {};
        descriptorPoolCreateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
        descriptorPoolCreateInfo.maxSets = 1;
        descriptorPoolCreateInfo.poolSizeCount = 1;
        descriptorPoolCreateInfo.pPoolSizes = &descriptorPoolSize;

        VK_CHECK_RESULT(vkCreateDescriptorPool(device, &descriptorPoolCreateInfo, NULL, &descriptorPool));

        // allocate single descriptor set
        VkDescriptorSetAllocateInfo descriptorSetAllocateInfo = {};
        descriptorSetAllocateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
        descriptorSetAllocateInfo.descriptorPool = descriptorPool;
        descriptorSetAllocateInfo.descriptorSetCount = 1;
        descriptorSetAllocateInfo.pSetLayouts = &descriptorSetLayout;

        VK_CHECK_RESULT(vkAllocateDescriptorSets(device, &descriptorSetAllocateInfo, &descriptorSet));
    }

    void VkOp::load_shader(const std::string& shaderfile)
    {
        uint32_t filelength;
        uint32_t* code = read_shader_file(filelength, shaderfile.c_str());
        VkShaderModuleCreateInfo createInfo = {};
        createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
        createInfo.pCode = code;
        createInfo.codeSize = filelength;
        VK_CHECK_RESULT(vkCreateShaderModule(device, &createInfo, NULL, &shaderModule));
        delete[] code;
    }

    void VkOp::create_pipeline(uint32_t pushconstants_size)
    {
        // shader stages
        VkPipelineShaderStageCreateInfo shaderStageCreateInfo = {};
        shaderStageCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        shaderStageCreateInfo.stage = VK_SHADER_STAGE_COMPUTE_BIT;
        shaderStageCreateInfo.module = shaderModule;
        shaderStageCreateInfo.pName = "main";

        // pipeline layout
        VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo = {};
        pipelineLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
        pipelineLayoutCreateInfo.setLayoutCount = 1;
        pipelineLayoutCreateInfo.pSetLayouts = &descriptorSetLayout;

        // push constants
        VkPushConstantRange pushConstantsRange = {};
        if (pushconstants_size > 0)
        {
            pushConstantsRange.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
            pushConstantsRange.offset = 0;
            pushConstantsRange.size = pushconstants_size;
            pipelineLayoutCreateInfo.pushConstantRangeCount = 1;
            pipelineLayoutCreateInfo.pPushConstantRanges = &pushConstantsRange;
        }
        else
        {
            pipelineLayoutCreateInfo.pushConstantRangeCount = 0;
            pipelineLayoutCreateInfo.pPushConstantRanges = nullptr;
        }
        
        VK_CHECK_RESULT(vkCreatePipelineLayout(device, &pipelineLayoutCreateInfo, NULL, &pipelineLayout));

        VkComputePipelineCreateInfo pipelineCreateInfo = {};
        pipelineCreateInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
        pipelineCreateInfo.stage = shaderStageCreateInfo;
        pipelineCreateInfo.layout = pipelineLayout;

        VK_CHECK_RESULT(vkCreateComputePipelines(device, NULL, 1, &pipelineCreateInfo, NULL, &pipeline));
    }
}
